diff --git a/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier.cc
--- a/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier.cc
+++ b/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -1558,28 +1558,37 @@ AlgebraicSimplifierVisitor::OptimizeDotO
       reshape->operand(0)->shape(), reshape->shape());
   CHECK_EQ(lhs_contracting_dims.size(), 1);
   if ((unmodified_dims.size() != reshape->shape().rank() - 1) ||
       absl::c_any_of(unmodified_dims, [&](const std::pair<int64, int64>& p) {
         return p.second == lhs_contracting_dims[0];
       })) {
     return nullptr;
   }
-  // Virtually pull the reshape into the dot. Now the dot is equivalent to a
-  // new dot with "unsquished" lhs contracting dims. We don't need to actually
-  // create a new dot instruction. We can just keep track of lhs and
-  // lhs_contracting_dims.
-  CHECK_GT(reshape->operand(0)->shape().rank(), reshape->shape().rank());
-  lhs_contracting_dims.Resize(
-      reshape->operand(0)->shape().rank() - reshape->shape().rank() + 1,
-      lhs_contracting_dims[0]);
-  absl::c_iota(lhs_contracting_dims, lhs_contracting_dims[0]);
+
+  // Virtually pull the reshape into the dot so the dot operates on the
+  // transpose, with "unsquished" lhs contracting dims.  The new contracting
+  // dims are all of the dims that are modified by the reshape -- that is, every
+  // dimension that's not in `unmodified_dims[i].first`.
+  //
+  // (We don't need to actually create a new dot instruction. We can just keep
+  // track of lhs and lhs_contracting_dims.)
+  absl::flat_hash_set<int64> unmodified_transpose_dims;
+  for (const auto& pair : unmodified_dims) {
+    unmodified_transpose_dims.insert(pair.first);
+  }
+  lhs_contracting_dims.Clear();
+  for (int64 i = 0; i < transpose->shape().dimensions_size(); ++i) {
+    if (!unmodified_transpose_dims.contains(i)) {
+      lhs_contracting_dims.Add(i);
+    }
+  }
   lhs = lhs->mutable_operand(0);
 
-  // Check if transpose only permutes the contracting dims.
+  // Check that the transpose only permutes the contracting dims.
   const auto& transpose_dims = transpose->dimensions();
   for (int64 i = 0; i < transpose_dims.size(); ++i) {
     if (transpose_dims[i] != i &&
         !absl::c_linear_search(lhs_contracting_dims, i)) {
       return nullptr;
     }
   }
   // Virtually pull the transpose into the dot. Now the dot is equivalent to
diff --git a/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
--- a/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
+++ b/google3/third_party/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
@@ -5297,10 +5297,57 @@ TEST_F(AlgebraicSimplifierTest, DotContr
       [](const Shape&, const Shape&) { return false; });
   options.set_is_layout_sensitive(true);
   // The transformation of moving transpose and reshape to the constant side is
   // layout insensitive. It should not happen if AlgebraicSimplifier is set up
   // to be layout sensitive.
   ASSERT_FALSE(AlgebraicSimplifier(options).Run(m.get()).ValueOrDie());
 }
 
+TEST_F(AlgebraicSimplifierTest, DotContractingReorder_SizeOneDimsNoChange) {
+  // This isn't transformed (notice that the relative order of the `2` and `3`
+  // dims doesn't change, so there's no opportunity here), but it's nonetheless
+  // an interesting testcase because of the presence of the size-1 dimensions.
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+     param = f32[1,2,5,3] parameter(0)
+     transpose = f32[1,5,2,3] transpose(param), dimensions={0,2,1,3}
+     reshape = f32[5,6] reshape(transpose)
+     constant = f32[6,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4}})
+     ROOT dot = f32[5,4] dot(reshape, constant),
+       lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie());
+}
+
+TEST_F(AlgebraicSimplifierTest, DotContractingReorder_SizeOneDims) {
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+     param = f32[1,2,3,5] parameter(0)
+     transpose = f32[1,3,2,5] transpose(param), dimensions={0,2,1,3}
+     reshape = f32[6,5] reshape(transpose)
+     constant = f32[6,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4}})
+     ROOT dot = f32[5,4] dot(reshape, constant),
+       lhs_contracting_dims={0}, rhs_contracting_dims={0}
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie());
+  auto shape1 = ShapeUtil::MakeShape(F32, {6, 5});
+  auto shape2 = ShapeUtil::MakeShape(F32, {1, 3, 2, 4});
+  auto shape3 = ShapeUtil::MakeShape(F32, {1, 2, 3, 4});
+  const HloInstruction* transpose;
+  ASSERT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Dot(
+                  m::Reshape(m::Parameter(0)).WithShapeCompatibleTo(&shape1),
+                  m::Reshape(m::Transpose(&transpose,
+                                          m::Reshape(m::Constant())
+                                              .WithShapeCompatibleTo(&shape2))
+                                 .WithShapeCompatibleTo(&shape3)))));
+  EXPECT_THAT(transpose->dimensions(), ElementsAre(0, 2, 1, 3));
+}
+
 }  // namespace
 }  // namespace xla

if (!is_iota(unmodified_transpose_dims)) { // I did not find an std or absl library for is_iota
  return nullptr;
}

  param = f32[2,5,1,3] parameter(0)
  transpose = f32[1,5,2,3] transpose(param), dimensions={2,1,0,3}
  reshape = f32[5,6] reshape(transpose)
  constant = f32[6,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4},{1,2,3,4}})
  ROOT dot = f32[5,4] dot(reshape, constant),
    lhs_contracting_dims={1}, rhs_contracting_dims={0}